# -*- coding: utf-8 -*-
# @Time    : 2023/9/29 10:36
# @Author  : Pan
# @Software: PyCharm
# @Project : VisualFramework
# @FileName: UNite
import math
import paddle
from paddle import nn
import paddle.nn.functional as F


def autopad(k, p=None):  # kernel, padding
    # Pad to 'same'
    if p is None:
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
    return p


def eval_act(act):
    if isinstance(act, str):
        return eval(act)
    elif act is True:
        return nn.Silu()
    elif isinstance(act, nn.Layer):
        return act
    else:
        return nn.Identity()


class AxialDWConv(nn.Layer):
    """ x` = x + DW_1xn(x) + DW_nx1(x)"""

    def __init__(self, in_channels, ksz=7, dsz=1) -> None:
        super(AxialDWConv, self).__init__()
        psz = autopad(ksz) * dsz
        self.dw1 = nn.Conv2D(in_channels, in_channels, kernel_size=(ksz, 1), padding=(psz, 0), dilation=dsz, groups=in_channels)
        self.dw2 = nn.Conv2D(in_channels, in_channels, kernel_size=(1, ksz), padding=(0, psz), dilation=dsz, groups=in_channels)
        self.mix = nn.Conv2D(in_channels, in_channels, kernel_size=1, padding=0)

    def forward(self, x):
        return self.mix(self.dw1(x) + self.dw2(x))


class DWConv(nn.Layer):
    def __init__(self, in_channels, mid_channels, out_channels, ksz=7, dsz=1):
        super(DWConv, self).__init__()
        psz = autopad(ksz) * dsz
        self.conv = nn.Sequential(
            nn.Conv2D(in_channels, mid_channels, 1),
            nn.Conv2D(mid_channels, mid_channels, ksz, dilation=dsz, padding=psz, groups=mid_channels),
            nn.Conv2D(mid_channels, out_channels, 1),
        )

    def forward(self, x):
        return self.conv(x)


class SELayer(nn.Layer):
    def __init__(self, num_channels, num_filters, reduction_ratio):
        super(SELayer, self).__init__()
        self.pool2d_gap = nn.AdaptiveAvgPool2D(1)
        self._num_channels = num_channels
        med_ch = int(num_channels / reduction_ratio)
        self.squeeze = nn.Linear(num_channels, med_ch)
        self.excitation = nn.Linear(med_ch, num_filters)

    def forward(self, input):
        pool = self.pool2d_gap(input)
        pool = paddle.squeeze(pool, axis=[2, 3])
        squeeze = self.squeeze(pool)
        squeeze = F.relu(squeeze)
        excitation = self.excitation(squeeze)
        excitation = F.sigmoid(excitation)
        excitation = paddle.unsqueeze(excitation, axis=[2, 3])
        out = input * excitation
        return out


class LayerConv(nn.Layer):
    """ The basic unit of ULitePan """

    def __init__(self, in_channels, out_channels, ksz=7, dsz=1, act=True, use_se=False) -> None:
        super().__init__()

        self.merge_conv = DWConv(in_channels, out_channels//4, out_channels, ksz)
        self.axial_conv = AxialDWConv(out_channels, ksz, dsz)

        self.act = eval_act(act)

        self.se = SELayer(out_channels, out_channels, 4) if use_se else None
        self.tail = nn.Conv2D(in_channels, out_channels, 1) if in_channels != out_channels else None

    def forward(self, x):
        y = self.merge_conv(x)
        y = self.axial_conv(y)
        y = self.act(y)
        if self.se is not None:
            y = self.se(y)
        if self.tail is not None:
            x = y + self.tail(x)
        else:
            x = x + y
        return x


class AxialDWBottleNeck(nn.Layer):
    """ BottleNeck implement by AxialDWConv """

    def __init__(self, in_channels, ksz=3, act=True, use_se=False) -> None:
        super().__init__()
        assert in_channels % 4 == 0
        mid_channels = in_channels // 4

        self.fore = nn.Conv2D(in_channels, mid_channels, kernel_size=(1, 1))
        self.adw1 = LayerConv(in_channels, mid_channels, ksz, 1, True)
        self.adw2 = LayerConv(in_channels, mid_channels, ksz, 2, True)
        self.adw3 = LayerConv(in_channels, mid_channels, ksz, 3, True)
        self.post = nn.Sequential(
            LayerConv(in_channels, in_channels, ksz, 1, True, use_se),
            eval_act(act)
        )

    def forward(self, x):
        x0 = self.fore(x)
        x1 = self.adw1(x)
        x2 = self.adw2(x)
        x3 = self.adw3(x)
        x4 = paddle.concat([x0, x1, x2, x3], axis=1)
        out = self.post(x4)
        return out


class Encoder(nn.Layer):
    def __init__(self, in_channels, out_channels, kernel_size, act, use_se):
        super(Encoder, self).__init__()
        self.encoder = LayerConv(in_channels, in_channels, kernel_size, 1, act, use_se)

        self.downsample = nn.Sequential(
            nn.Conv2D(in_channels=in_channels, out_channels=in_channels, kernel_size=2, stride=2, groups=in_channels),
            nn.Conv2D(in_channels, out_channels, kernel_size=1),
            eval_act(act)
        )

    def forward(self, x):
        x = x + self.encoder(x)
        y = self.downsample(x)
        return x, y


class Decoder(nn.Layer):
    def __init__(self, in_channels, out_channels, kernel_size, act, use_se):
        super(Decoder, self).__init__()
        self.upsample = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            LayerConv(in_channels, out_channels, kernel_size, 1, act, use_se)
        )

        self.encoder = LayerConv(out_channels * 2, out_channels, kernel_size, 1, act, use_se)

    def forward(self, x, y):
        y = self.upsample(y)
        x = self.encoder(paddle.concat([x, y], axis=1))
        return x


class ULitePan(nn.Layer):
    """ Implementation of "1M PARAMETERS ARE ENOUGH? A LIGHTWEIGHT CNN-BASED MODEL FOR MEDICAL IMAGE SEGMENTATION" """

    def __init__(self, configs) -> None:
        super(ULitePan, self).__init__()
        in_channels = configs["in_channels"] if "in_channels" in configs.keys() else 3
        encode_dim = configs["encode_dim"] if "encode_dim" in configs.keys() else 8
        mid_channels = configs["mid_channels"] if "mid_channels" in configs.keys() else (16, 32, 64, 128, 256, 512)
        num_classes = configs["num_classes"] if "num_classes" in configs.keys() else 3
        act = configs["act"] if "act" in configs.keys() else True
        use_se = configs["use_se"] if "use_se" in configs.keys() else True
        kernel_size = 7

        # 对图像做编码和大尺度的采集
        self.inconv = DWConv(in_channels, encode_dim, mid_channels[0], 7)

        # 分成n阶段【n=5】
        self.nc = len(mid_channels) - 1
        # 编码模块
        self.encoders = nn.LayerList([Encoder(mid_channels[i], mid_channels[i + 1], kernel_size, act, use_se) for i in range(self.nc)])
        # 对最深层的特征做融合
        self.bottle = AxialDWBottleNeck(mid_channels[self.nc], 5, act, use_se)

        # 解码模块
        self.decoders = nn.LayerList([Decoder(mid_channels[i], mid_channels[i-1], kernel_size, act, use_se) for i in range(self.nc, 0, -1)])

        # 输出
        self.outconv = nn.Conv2D(mid_channels[0], num_classes, kernel_size=(1, 1))

    def forward(self, x):
        y = self.inconv(x)

        feats = []
        for i in range(len(self.encoders)):
            x, y = self.encoders[i](y)
            feats.append(x)

        y = self.bottle(y)
        for i in range(len(self.decoders)):
            y = self.decoders[i](feats[len(feats)-i-1], y)

        out = self.outconv(y)
        return [out]


if __name__ == '__main__':
    model = ULitePan({
        "mid_channels": (12, 24, 48, 96, 192, 384)
    })
    paddle.flops(model, [1, 3, 512, 512])
